from langchain_community.vectorstores import Chroma
from model.embeddings import DashScopeEmbeddings
import os


def initialize_vectorstore(documents, embedding_model="text-embedding-v4", dimensions=1024):
    if not documents:
        print("没有可处理的文档")
        return None

    texts = [doc.page_content for doc in documents]
    metadatas = [doc.metadata for doc in documents]

    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        print("未设置API密钥，无法生成嵌入")
        return None

    embeddings = DashScopeEmbeddings(
        api_key=api_key,
        model=embedding_model,
        dimensions=dimensions
    )

    # 批量处理文档（避免单次请求过大）
    batch_size = 10
    vectorstore = None

    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i:i + batch_size]
        batch_metadatas = metadatas[i:i + batch_size]

        if vectorstore is None:
            # 首次创建向量存储
            vectorstore = Chroma.from_texts(
                texts=batch_texts,
                embedding=embeddings,
                metadatas=batch_metadatas,
                persist_directory="./chroma_db"
            )
        else:
            # 后续批次添加到现有存储
            vectorstore.add_texts(
                texts=batch_texts,
                metadatas=batch_metadatas
            )
    return vectorstore


def load_existing_vectorstore(embedding_model="text-embedding-v4", dimensions=1024):
    """加载已存在的向量存储"""
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        print("未设置API密钥，无法加载向量存储")
        return None

    embeddings = DashScopeEmbeddings(
        api_key=api_key,
        model=embedding_model,
        dimensions=dimensions
    )

    # 检查是否存在现有向量库
    if os.path.exists("./chroma_db") and os.listdir("./chroma_db"):
        try:
            return Chroma(
                persist_directory="./chroma_db",
                embedding_function=embeddings
            )
        except Exception as e:
            print(f"加载现有向量存储失败: {e}")
            return None
    return None
